/*
    Author: @5mukx
*/

use core::arch::asm;
use native::{LoadLibraryA, MessageBoxA, LDR_DATA_TABLE_ENTRY, LIST_ENTRY, PEB};
use winapi::shared::ntdef::LPCSTR;
use winapi::um::winuser::MB_OK;
use winapi::shared::minwindef::{DWORD, HINSTANCE, FARPROC};
use winapi::shared::basetsd::{UINT_PTR, DWORD64};
use winapi::um::winnt::{
    IMAGE_DOS_HEADER, IMAGE_EXPORT_DIRECTORY,
    IMAGE_NT_HEADERS64,
};
use winapi::shared::windef::HWND;
use std::os::raw::c_char;

pub mod native;

#[inline]
#[cfg(target_pointer_width = "64")]
unsafe fn __readgsqword(offset: DWORD) -> DWORD64 {
    let out: u64;
    asm!(
        "mov {}, gs:[{:e}]",
        lateout(reg) out,
        in(reg) offset,
        options(nostack, pure, readonly),
    );
    out
}

unsafe fn get_module_base_addr(module_name: &str) -> HINSTANCE{
    let peb_offset = __readgsqword(0x60) as *const u64;
    let rf_peb: *const PEB = peb_offset as *const PEB;

    // borrow it insted of moving !
    let peb = &*rf_peb;

    let mut p_ldr_data_table_entry: *const LDR_DATA_TABLE_ENTRY = 
        (*peb.Ldr).InMemoryOrderModuleList.Flink as *const LDR_DATA_TABLE_ENTRY;

    let mut p_list_entry = &(*peb.Ldr).InMemoryOrderModuleList as *const LIST_ENTRY;

    loop{
        let buffer = std::slice::from_raw_parts(
            (*p_ldr_data_table_entry).FullDllName.Buffer, 
            (*p_ldr_data_table_entry).FullDllName.Length as usize / 2,
        );

        let dll_name = String::from_utf16_lossy(buffer);
        
        if dll_name.to_lowercase().starts_with(module_name){
            let module_base: HINSTANCE = (*p_ldr_data_table_entry).Reserved2[0] as HINSTANCE;
            return module_base;
        }

        if p_list_entry == (*peb.Ldr).InMemoryOrderModuleList.Blink {
            println!("Module not found!");
            return 0 as HINSTANCE;
        }

        p_list_entry = (*p_list_entry).Flink;
        p_ldr_data_table_entry = (*p_list_entry).Flink as *const LDR_DATA_TABLE_ENTRY;
    }
}

// -> FARPROC
fn get_proc_addr(module_handle: HINSTANCE, function_name: &str) -> FARPROC{
    
    unsafe {
        let dos_headers = module_handle as *const IMAGE_DOS_HEADER;
        let nt_headers = (module_handle as UINT_PTR + (*dos_headers).e_lfanew as UINT_PTR) 
            as *const IMAGE_NT_HEADERS64;

        let data_directory = &(*nt_headers).OptionalHeader.DataDirectory[0];
        let export_directory = (module_handle as UINT_PTR + (*data_directory).VirtualAddress as UINT_PTR) 
            as *const IMAGE_EXPORT_DIRECTORY;

        let mut address_array = module_handle as UINT_PTR + (*export_directory).AddressOfFunctions as UINT_PTR;
        let mut name_array = module_handle as UINT_PTR + (*export_directory).AddressOfNames as UINT_PTR;
        let mut name_ordinals = module_handle as UINT_PTR + (*export_directory).AddressOfNameOrdinals as UINT_PTR;

        loop {

            let name_offset: u32 = *(name_array as *const u32);
            let current_function_name = std::ffi::CStr::from_ptr(
                (module_handle as UINT_PTR + name_offset as UINT_PTR) as *const c_char
            ).to_str().unwrap();

            if current_function_name == function_name{
                address_array = address_array + (*(name_ordinals as *const u16) as UINT_PTR * std::mem::size_of::<DWORD>() as UINT_PTR);
                let fun_addr: FARPROC = std::mem::transmute(module_handle as UINT_PTR + *(address_array as *const u32) as UINT_PTR);
                return fun_addr;
            }

            name_array = name_array + std::mem::size_of::<DWORD>() as UINT_PTR;
            name_ordinals = name_ordinals + std::mem::size_of::<u16>() as UINT_PTR;
        }
    }
}

fn main(){

    unsafe{
        println!("[+] Getting base address of kernel32.dll");
        let kernel32_base_address: HINSTANCE = get_module_base_addr("kernel32.dll");

        println!("[+] Dynamically resolving LoadLibraryA");

        let load_library_a: LoadLibraryA = std::mem::transmute(get_proc_addr(kernel32_base_address, "LoadLibraryA"));
        
        println!("[*] LoadLibraryA Address: {:?}", load_library_a);


        println!("[+] Load user32.dll");
        load_library_a("user32.dll\0".as_ptr() as LPCSTR);

        println!("[+] Getting base address of user32.dll");

        let user32_base_address: HINSTANCE = get_module_base_addr("user32.dll");

        println!("[*] user32.dll Address: {:?}", user32_base_address);

        println!("[+] Dynamically resolve MessageBoxA");


        let message_box_a: MessageBoxA = 
            std::mem::transmute(get_proc_addr(user32_base_address, "MessageBoxA"));

        println!("[*] MessageBoxA Address: {:?}", message_box_a);

        message_box_a(
            0 as HWND,
            "Hey, 5mukx here. We Resolved it Dynamically.\0".as_ptr() as LPCSTR,
            "MessageBoxA\0".as_ptr() as LPCSTR,
            MB_OK
        );
    }   
}
